import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from d2l import torch as d2l

from utils.file_download_util import FileDownloader


class ModelTrainer:
    @staticmethod
    def log_rmse(net, features, labels):
        clipped_preds = torch.clamp(net(features), 1, float('inf'))
        rmse = torch.sqrt(loss(torch.log(clipped_preds), torch.log(labels)))
        return rmse.item()

    @staticmethod
    def get_default_net(in_features):
        net = nn.Sequential(nn.Linear(in_features, 1))
        return net

    @staticmethod
    def train(net=None, train_features=None, train_labels=None, test_features=None, test_labels=None,
              num_epochs=100, learning_rate=0.01, weight_decay=0, batch_size=64):
        if net is None:
            net = ModelTrainer.get_default_net(train_features.shape[1])
        train_ls, test_ls = [], []
        train_iter = d2l.load_array((train_features, train_labels), batch_size)
        optimizer = optim.Adam(net.parameters(), lr=learning_rate, weight_decay=weight_decay)

        for epoch in range(num_epochs):
            for X, y in train_iter:
                optimizer.zero_grad()
                l = loss(net(X), y)
                l.backward()
                optimizer.step()
            train_ls.append(ModelTrainer.log_rmse(net, train_features, train_labels))
            if test_labels is not None:
                test_ls.append(ModelTrainer.log_rm.pyrmse(net, test_features, test_labels))
        return train_ls, test_ls


# 文件下载示例
if __name__ == '__main__':
    # 1-定义要下载的文件
    DATA_URL = 'http://d2l-data.s3-accelerate.amazonaws.com/'
    FileDownloader.DATA_HUB = {
        # 二元组分别对应其url和密钥
        'kaggle_house_train': (DATA_URL + 'kaggle_house_pred_train.csv', '585e9cc93e70b39160e7921475f9bcd7d31219ce'),
        'kaggle_house_test': (DATA_URL + 'kaggle_house_pred_test.csv', 'fa19780a7b011d9b009e8bff8e99922a8ee2eb90'),
    }

    # 2-进行文件下载-查看数据格式
    train_data = pd.read_csv(FileDownloader.download('kaggle_house_train'))
    test_data = pd.read_csv(FileDownloader.download('kaggle_house_test'))
    print('----------【工具】train_data.shape：', train_data.shape)
    print('----------【工具】test_data.shape：', test_data.shape)

    # 3-查看数据格式
    print('----------train_data.shape：', train_data.shape)
    print('----------test_data.shape：', test_data.shape)

    # 4-查看前四个样本的前四个后最后两个特征，以及相应的标价
    print('----------train_data.iloc：', train_data.iloc[0:4, [0, 1, 2, 3, -3, -2, -1]])

    # 5-【数据预处理】我们将其id信息从数据集中删除
    # 去掉ID后的所有样本,都表示对于训练集/测试集的所有样本，从下标1一直到最后，使用pd的concat将二者合并
    all_features = pd.concat((train_data.iloc[:, 1:-1], test_data.iloc[:, 1:]))
    print('----------all_features.shape：', all_features.shape)

    # 6-【数据预处理-数值特征】将所有缺失值替换为相应特征的平均值。然后进行feature scaling（先标准化再替换会更方便，因为标准化后均值就是0了）
    # 假设不是'object'类的数据就是数值类，把所有数值特征提取出来
    numeric_feaures = all_features.select_dtypes(include=['number']).columns
    # 对数值特征进行标准化(实际情况拿不到测试集的时候可以用训练集的均值方差代替)
    all_features[numeric_feaures] = all_features[numeric_feaures].apply(
        lambda x: (x - x.mean()) / (x.std()))
    # 把NAN变成均值，不建议改成0
    all_features[numeric_feaures] = all_features[numeric_feaures].fillna(all_features[numeric_feaures].mean())

    # 7-【数据预处理-离散值特征】进行one-hot独热编码替换
    all_features = pd.get_dummies(all_features, dummy_na=True)
    # 将布尔值转换为0和1(当前Python版本默认使用True/False)
    all_features = all_features.astype(int)
    # 填充缺失值为相应列的均值
    all_features = all_features.fillna(all_features.mean())

    # 8-【数据预处理-分割数据集】将数据集分割为训练集和测试集
    # 转换为torch张量
    n_train = train_data.shape[0]
    # 从all_features中取出前n_train行作为训练集，后n_test行作为测试集
    train_features = torch.tensor(all_features[:n_train].values, dtype=torch.float32)
    test_features = torch.tensor(all_features[n_train:].values, dtype=torch.float32)
    # 从train_data中取出前n_train行作为训练集，后n_test行作为测试集
    train_labels = torch.tensor(train_data.SalePrice.values, dtype=torch.float32).view(-1, 1)

    # 9-【数据预处理-查看数据】查看数据
    print('----------train_features.shape：', train_features.shape)
    print('----------train_labels.shape：', train_labels.shape)
    print('----------test_features.shape：', test_features.shape)

    # 10-【模型训练】训练⼀个带有损失平⽅的线性模型
    loss = nn.MSELoss()
    in_features = train_features.shape[-1]  # 提取所有训练特征

    # 11-【模型训练】进行模型训练
    # k折交叉验证的折数，训练轮数，学习率，权重衰减，批量大小-设置超参数
    k, num_epochs, lr, weight_decay, batch_size = 5, 100, 5, 0, 64
    # 创建默认的网络模型
    net = ModelTrainer.get_default_net(in_features)
    # 调用k_fold函数进行训练和验证，返回训练损失和验证损失
    train_l, valid_l = ModelTrainer.train(net, train_features, train_labels, num_epochs, lr, weight_decay, batch_size)
    print('直接交叉验证全部搞定！')
    print(f'直接交叉验证：平均训练log rmse：{train_l},'
          f'平均验证log rmse：{valid_l}')
